import wenxin_api
from config import AI_CONFIG
from wenxin_api.tasks.free_qa import FreeQA


class WenXinBot:

    def __init__(self):
        config = AI_CONFIG["wenxin"]
        wenxin_api.ak = config["ak"]
        wenxin_api.sk = config["sk"]
        self.task_prompt = config["task_prompt"]
        self.seq_len = config["seq_len"]
        self.topp = config["topp"]
        self.typeId = config["typeId"]

    def ask(self, prompt):

        input_dict = {
            # 用户输入文本 string [1, 1000]
            "text": prompt,
            # 最大生成长度 int [1, 1000]
            "seq_len": self.seq_len,
            # 多样性 float [0.0,1.0]
            "topp": self.topp,
            # 任务类型
            "task_prompt": self.task_prompt
        }

        response = FreeQA.create(**input_dict)
        return response.text

# 参数说明页面https://wenxin.baidu.com/wenxin/docs#dl6tgxw5f
# 【task_prompt】 指定预置的任务模板，效果更好。
# PARAGRAPH：引导模型生成一段文章；
# SENT：引导模型生成一句话；
# ENTITY：引导模型生成词组；
# Summarization：摘要；
# MT：翻译；
# Text2Annotation：抽取；
# Correction：纠错；
# QA_MRC：阅读理解；
# Dialogue：对话；
# QA_Closed_book: 闭卷问答；
# QA_Multi_Choice：多选问答；
# QuestionGeneration：问题生成；
# Paraphrasing：复述；
# NLI：文本蕴含识别；
# SemanticMatching：匹配；
# Text2SQL：文本描述转SQL；
# TextClassification：文本分类；
# SentimentClassification：情感分析；
# zuowen：写作文；
# adtext：写文案；
# couplet：对对联；
# novel：写小说；
# cloze：文本补全；
# Misc：其它任务。